import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from torchvision.datasets.folder import default_loader
from PIL import Image
from torchvision.datasets import ImageFolder

data_transforms = {
    'train':
        transforms.Compose([
            transforms.Resize(64),
            transforms.RandomRotation(45),  # 随机旋转，-45到45度之间随机选
            transforms.CenterCrop(64),  # 从中心开始裁剪
            transforms.RandomHorizontalFlip(p=0.5),  # 随机水平翻转 选择一个概率概率
            transforms.RandomVerticalFlip(p=0.5),  # 随机垂直翻转
            transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),
            # 参数1为亮度，参数2为对比度，参数3为饱和度，参数4为色相
            transforms.RandomGrayscale(p=0.025),  # 概率转换成灰度率，3通道就是R=G=B
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 均值，标准差
        ]),
    'valid':
        transforms.Compose([
            transforms.Resize(64),
            transforms.CenterCrop(64),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
}


class FlowerDataset(Dataset):
    def __init__(self, data_dir, anno_file, transform=None):
        super(FlowerDataset, self).__init__()
        self.data_dir = data_dir
        self.anno_file = anno_file
        self.img_label = self.load_annotations()
        self.img = [os.path.join(data_dir, img) for img in list(self.img_label.keys())]
        self.label = [label for label in list(self.img_label.values())]
        self.transform = transform

    def __len__(self):
        return len(self.img)

    def __getitem__(self, index):
        image = default_loader(self.img[index])
        # print("default_loader:", type(image))
        # image = Image.open(self.img[index])
        # print(type(image))
        label = self.label[index]
        if self.transform:
            image = self.transform(image)
        # print(type(image))
        label = torch.from_numpy(label)
        return image, label

    def load_annotations(self):
        data_infos = {}
        with open(self.anno_file) as anno:
            samples = [sam.strip().split(' ') for sam in anno.readlines()]
        for file_name, gt_label in samples:
            data_infos[file_name] = np.array(gt_label, dtype=np.int64)
        return data_infos


if __name__ == "__main__":
    train_path = "/home/stark/algo-env/datasets/flower_data/train_filelist"
    val_path = "/home/stark/algo-env/datasets/flower_data/val_filelist"
    train_file = "/home/stark/algo-env/datasets/flower_data/train.txt"
    val_file = "/home/stark/algo-env/datasets/flower_data/val.txt"

    train_dataset = FlowerDataset(train_path, train_file, transform=data_transforms["train"])
    val_dataset = FlowerDataset(val_path, val_file, transform=data_transforms["valid"])
    print(train_dataset.img_label)
    print(train_dataset.label)
    # train_loader = DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)
    # img, label = next(iter(train_loader))
    # sample = img[0].squeeze()
    # sample = sample.permute(1, 2, 0).numpy()
    # sample *= [0.229, 0.224, 0.225]
    # sample += [0.485, 0.456, 0.406]
    # plt.imshow(sample)
    # plt.show()
    # print("Label is : {}".format(label[0].numpy()))
    # for img, label in train_loader:
    #     print(img, label)
    # flower_dataset = ImageFolder(r"/home/stark/algo-env/datasets/flower_data/train")
    # print(flower_dataset.imgs)
    # print(flower_dataset.class_to_idx)
